package com.lhz.collector.source;

import com.google.common.collect.Lists;
import com.lhz.collector.common.Constant;
import org.apache.flume.Context;
import org.apache.flume.Event;
import org.apache.flume.EventDrivenSource;
import org.apache.flume.conf.Configurable;
import org.apache.flume.conf.Configurables;
import org.apache.flume.instrumentation.SourceCounter;
import org.apache.flume.source.AbstractSource;
import org.jboss.netty.bootstrap.ServerBootstrap;
import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.channel.*;
import org.jboss.netty.channel.socket.nio.NioServerSocketChannelFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
public class NettyTcpSource extends AbstractSource
                             implements EventDrivenSource, Configurable {
  private static final Logger logger = LoggerFactory.getLogger(NettyTcpSource.class);

  private int port;
  private String host = null;
  private Integer batchSize;
  private Channel nettyChannel;
  private Integer eventSize;
  private SourceCounter sourceCounter;

  public class syslogTcpHandler extends SimpleChannelHandler {

    private NettyUtils syslogUtils = new NettyUtils();

    public void setEventSize(int eventSize) {
      syslogUtils.setEventSize(eventSize);
    }

    @Override
    public void messageReceived(ChannelHandlerContext ctx, MessageEvent mEvent) {
      List<Event> events = Lists.newArrayList();
      ChannelBuffer buff = (ChannelBuffer) mEvent.getMessage();
      String IP=((InetSocketAddress)mEvent.getRemoteAddress()).getAddress().getHostAddress();
      while (buff.readable()) {
        Event e = syslogUtils.extractEvent(buff);
        if (e == null) {
          logger.debug("Parsed partial event, event will be generated when rest of the event is received.");
          continue;
        }
        sourceCounter.incrementEventReceivedCount();
        try {
          e.getHeaders().put(Constant.IP,IP);
          events.add(e);
          sourceCounter.incrementEventAcceptedCount();
          if(events.size()>=batchSize){
            getChannelProcessor().processEventBatch(events);
            events.clear();
          }
        } catch (Exception ex) {
          logger.error("Error writting to channel, event dropped", ex);
        }
      }
      if(events.size()>0){
        getChannelProcessor().processEventBatch(events);
      }
    }
  }

  @Override
  public void start() {
    ChannelFactory factory = new NioServerSocketChannelFactory(
        Executors.newCachedThreadPool(), Executors.newCachedThreadPool());

    ServerBootstrap serverBootstrap = new ServerBootstrap(factory);
    serverBootstrap.setPipelineFactory(new ChannelPipelineFactory() {
      @Override
      public ChannelPipeline getPipeline() {
        syslogTcpHandler handler = new syslogTcpHandler();
        handler.setEventSize(eventSize);
        return Channels.pipeline(handler);
      }
    });

    logger.info("Syslog TCP Source starting...");

    if (host == null) {
      nettyChannel = serverBootstrap.bind(new InetSocketAddress(port));
    } else {
      nettyChannel = serverBootstrap.bind(new InetSocketAddress(host, port));
    }

    sourceCounter.start();
    super.start();
  }

  @Override
  public void stop() {
    logger.info("Syslog TCP Source stopping...");
    logger.info("Metrics: {}", sourceCounter);

    if (nettyChannel != null) {
      nettyChannel.close();
      try {
        nettyChannel.getCloseFuture().await(60, TimeUnit.SECONDS);
      } catch (InterruptedException e) {
        logger.warn("netty server stop interrupted", e);
      } finally {
        nettyChannel = null;
      }
    }

    sourceCounter.stop();
    super.stop();
  }

  @Override
  public void configure(Context context) {
    Configurables.ensureRequiredNonNull(context,"port");
    port = context.getInteger("port");
    host = context.getString("host");
    batchSize=context.getInteger("batchSize",100);
    eventSize = context.getInteger("eventSize", NettyUtils.DEFAULT_SIZE);
    if (sourceCounter == null) {
      sourceCounter = new SourceCounter(getName());
    }
  }
}
